#############################################################################################
#############################################################################################
# 
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD, 
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be 
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or 
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of 
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this 
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are 
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software, 
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE 
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME 
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS. 
#############################################################################################
#############################################################################################

import sys, os

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')
    
# Import Data Assimilation class
try:
    from cdar.assimilation.assimilation_var import DA_var
except:
    raise ImportError('The DA_var library cannot be found!')

# Import Analysis class
try:
    from cdar.analysis_generator.var.var import VAR
except:
    raise ImportError('The VAR library cannot be found!')
    
# Import Model class
try:
    from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
except:
    raise ImportError('The model library cannot be found!')

class VAR_DA(DA_var):
    
    # Attribute to read a first guess from a file. 
    def fg_read(self,ParamFileName):
        DA_var.fg_read(self,ParamFileName)
        self.F_U,self.F_V,self.F_H = np.split(self.fg,self.domain.VarDim)
    
    def obs_read(self,DataFilePath):
                # read the observation files available during the assimilation window at each history_interval
                DA_var.obs_read(self,DataFilePath)
                
                # Get the statistics of the observational data
                self.u_obs_num = 0
                self.v_obs_num = 0
                self.h_obs_num = 0
                for i in range(len(self.innov_time)): 
                    print 'In observation file at %s' % str(self.innov_time[i])
                    [obs_u, obs_v, obs_h] = np.split(self.obs_list[i], self.domain.VarDim)
                    print '    The number of U observations is %d' % obs_u[~np.isnan(obs_u)].size
                    self.u_obs_num += obs_u[~np.isnan(obs_u)].size
                    print '    The number of V observations is %d' % obs_v[~np.isnan(obs_v)].size
                    self.v_obs_num += obs_v[~np.isnan(obs_v)].size
                    print '    The number of H observations is %d' % obs_h[~np.isnan(obs_h)].size
                    self.h_obs_num += obs_h[~np.isnan(obs_h)].size
                self.total_obs_num = self.u_obs_num + self.v_obs_num + self.h_obs_num


def VAR4D():

    try:
        import ESMF
    except:
        raise ImportError('The ESMP library cannot be found!')
        
    try:
        from cdar.utilities import namelist_read as nml
    except:
        raise ImportError('The nml library cannot be found!')

    EnSize = 1
    DataFilePath='./data'  # Path to store or read FG, B and OBS
    
    # start up ESMF
    # this call is not necessary unless you want to to override the
    # default options:
    #  LogKind = NONE
    #  debug = False
    manager = ESMF.Manager(logkind=ESMF.LogKind.MULTI, debug=True)
    
    # inquire for rank and proc from ESMF Virtual Machine
    localPet = manager.local_pet
    petCount = manager.pet_count
    print "localPet = %d and petCount = %d " % (localPet, petCount)
    
    # opening remarks
    if localPet == 0:
        print "Welcome to the Variational Data Assimilation !"
            
    # Specify analysis method
    analysis_method = VAR(optimization_algrm_choice='L-BFGS-B')
    print analysis_method

    # Read in parameters for model configurations
    params = nml.namelist_read(DataFilePath)

    # Set up the model
    domain = model(params)

    # Print domain info 
    if localPet == 0:
        print domain  
    
    # Redefine the history output interval
    # For 4D-Var, it has to be the timestep.
    domain.history_interval = domain.timestep
    domain.model_output_history()
    
    # Specify data assimilation method
    DA_method = VAR_DA(domain,DataFilePath,EnSize,analysis_method)

    # Read/Write initialization/parametrization file to correct place.
    
    # Read FG
    fg_file = os.path.join(DataFilePath,'FG_'+domain.Name+'_'+params['start_time']+'.npz')
    print "Reading the FG in file: %s" % fg_file
    DA_method.fg_read(fg_file)
    
    # Strip off the boundaries
    DA_method.fg = domain.strip_off_boundaries(np.vstack((DA_method.F_U,DA_method.F_V,DA_method.F_H)))
    
    # Load the observations
    DA_method.obs_read(DataFilePath)

    # Load the B
    B_file = os.path.join(DataFilePath,'B_'+domain.Name+'_'+'1000'+'_'+str(domain.nx)+'_'+params['start_time']+'.npz')
    print "Reading the BE in file: %s" %  B_file
    DA_method.B_read(B_file)

    # Run data assimilation routine
    Analysis = DA_method.DArun()
    
    # Save analysis
    Ana_file = os.path.join(DataFilePath,'Analysis_'+domain.Name+'_'+params['start_time']+'.npz')
    print "Writting the Analysis in file: %s" %  Ana_file
    DA_method.analysis_write(Analysis, Ana_file)
    
if __name__ == '__main__':
    sys.exit(VAR4D())
